# if not utility.has_collection(DOC_COLLECTION_NAME):
from pymilvus import (
    connections,
    FieldSchema,
    DataType,
    CollectionSchema,
    Collection,
    utility,
)

connections.connect(
    "default", host="192.168.1.62", port="19530", db_name="zhangrenyang"
)
# 定义集合的字段结构
fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="embedding1", dtype=DataType.FLOAT_VECTOR, dim=128),
    FieldSchema(name="embedding2", dtype=DataType.FLOAT_VECTOR, dim=128),
]
# 定义集合的名称
collection_name = "blog"
# 创建集合的模式对象
schema = CollectionSchema(fields, description=collection_name)

try:
    collection = None
    if utility.has_collection(collection_name):
        # 获取集合
        collection = Collection(collection_name)
    else:
        # 创建新的集合
        collection = Collection(collection_name, schema)
        print("集合创建成功")
    # 定义索引参数配置
    index_params = {
        "metric_type": "COSINE",  # L2 COSINE IP
        "index_type": "IVF_FLAT",  # 索引类型 倒排索引+暴力遍历
        "params": {"nlist": 128},  # nlist 越大招回率越高速度越慢
    }
    collection.create_index("embedding", index_params)
    print("索引创建成功")
    # 将集合加载到内存中以支持搜索操作
    collection.load()
    print("集合已经加载到内存中")
except Exception as e:
    print(e)
